import { StateGraph } from "@langchain/langgraph";
import { ChatOpenAI } from "@langchain/openai";
import { BaseMessage, HumanMessage, ToolMessage, SystemMessage } from "@langchain/core/messages";
import { StructuredTool, DynamicStructuredTool } from "@langchain/core/tools";
import * as http from 'http';
import { z } from 'zod';

import { OPENAI_API_KEY, AGENT_MODEL, OPENAI_BASE_URL, MCP_SERVER_PORT, DECISION_MODEL } from "./config";

type AgentState = {
    messages: BaseMessage[];
};

export class SpreadJSAgent {
    private app: StateGraph<AgentState>;
    private model: ChatOpenAI;
    private tools: StructuredTool[];

    constructor() {
        // Initialize without tools first, they will be loaded asynchronously
        this.tools = [];
        this.model = new ChatOpenAI({
            apiKey: OPENAI_API_KEY,
            modelName: AGENT_MODEL,
            configuration: {
                baseURL: OPENAI_BASE_URL
            },
            temperature: 0.3,  // Increase temperature for more conversational responses
            streaming: true    // Enable streaming for better UX
        });
        this.app = this.buildGraph();
    }

    // Initialize tools by fetching from MCP Server
    public async initialize(): Promise<void> {
        try {
            this.tools = await this.fetchMCPTools();
            this.model = this.model.bindTools(this.tools) as ChatOpenAI;
            console.log('[Agent] Initialized with MCP tools:', this.tools.map(t => t.name));
        } catch (error) {
            console.error('[Agent] Failed to initialize MCP tools:', error);
            throw error;
        }
    }

    // Fetch tools from MCP Server via HTTP
    private async fetchMCPTools(): Promise<StructuredTool[]> {
        const response = await this.httpRequest('GET', '/tools');
        const mcpTools = response.tools;

        return mcpTools.map((mcpTool: any) => {
            const schema = this.buildZodSchema(mcpTool.schemaDescription);

            return new DynamicStructuredTool({
                name: mcpTool.name,
                description: mcpTool.description,
                schema: schema as any, // Explicitly cast to avoid deep type inference
                func: async (args: any) => {
                    // 所有工具都使用统一的MCP调用，不再特殊处理execute-spreadjs
                    return await this.executeMCPTool(mcpTool.name, args);
                },
            });
        });
    }

    // Build proper zod schema from schema description
    private buildZodSchema(schemaDescription: any): z.ZodObject<any> {
        if (!schemaDescription || schemaDescription.type !== 'object') {
            return z.object({});
        }

        // Handle specific known schemas to avoid complex type inference
        const properties = schemaDescription.properties || {};

        if (properties.keyword) {
            // api-doc-search tool schema
            return z.object({
                keyword: z.string().describe(properties.keyword.description || 'Search keyword')
            });
        } else if (properties.execute_logic && properties.query_logic) {
            // execute-spreadjs tool schema
            return z.object({
                execute_logic: z.string().describe(properties.execute_logic.description || 'Execute logic'),
                query_logic: z.string().describe(properties.query_logic.description || 'Query logic')
            });
        }

        // Fallback to generic object schema
        return z.object({});
    }

    // Execute tool via HTTP request to MCP Server
    private async executeMCPTool(toolName: string, args: any): Promise<any> {
        const response = await this.httpRequest('POST', '/execute', {
            toolName: toolName,
            args: args
        });
        return response.result;
    }

    // HTTP request helper
    private httpRequest(method: string, path: string, data?: any): Promise<any> {
        return new Promise((resolve, reject) => {
            const postData = data ? JSON.stringify(data) : undefined;

            const options = {
                hostname: 'localhost',
                port: MCP_SERVER_PORT,
                path: path,
                method: method,
                headers: {
                    'Content-Type': 'application/json',
                    ...(postData && { 'Content-Length': Buffer.byteLength(postData) })
                }
            };

            const req = http.request(options, (res) => {
                let body = '';
                res.on('data', chunk => {
                    body += chunk;
                });

                res.on('end', () => {
                    try {
                        const response = JSON.parse(body);
                        if (response.success) {
                            resolve(response);
                        } else {
                            reject(new Error(response.error || 'Unknown error'));
                        }
                    } catch (error) {
                        reject(new Error('Failed to parse response: ' + body));
                    }
                });
            });

            req.on('error', (error) => {
                reject(new Error(`HTTP request failed: ${error.message}`));
            });

            if (postData) {
                req.write(postData);
            }

            req.end();
        });
    }

    private buildGraph(): StateGraph<AgentState> {
        const workflow = new StateGraph<AgentState>({
            channels: {
                messages: {
                    value: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y),
                    default: () => [],
                }
            },
        });

        // 节点 1: action - 执行工具调用和LLM推理
        workflow.addNode("action", this.actionNode.bind(this));
        // 节点 2: check - AI自主决策下一步行动
        workflow.addNode("check", this.checkNode.bind(this));

        (workflow as any).setEntryPoint("action");

        // action 执行后总是去 check 进行决策
        (workflow as any).addEdge("action", "check");

        // check 节点进行AI自主决策
        (workflow as any).addConditionalEdges("check", async (state: AgentState) => {
            const decision = await this.aiDecideNextAction(state);
            console.log(`[Workflow] AI自主决策: ${decision}`);

            if (decision === "__end__") {
                return "__end__";
            } else {
                return "action"; // CONTINUE_SEARCH 或 EXECUTE_CODE 都回到action执行
            }
        });

        return workflow;
    }

    // AI自主决策下一步行动 - 官方正确方式：让AI基于完整上下文自主决策
    private async aiDecideNextAction(state: AgentState): Promise<string> {
        // 使用官方方式：让AI基于完整的消息历史进行决策，而不是手动解析
        const decisionPrompt = `
基于以上完整的对话历史，决定下一步最合适的行动：

1. CONTINUE_SEARCH - 如果需要更多SpreadJS技术信息
2. EXECUTE_CODE - 如果有足够信息可以生成SpreadJS代码
3. TASK_COMPLETE - 如果用户请求已完全满足（请先向用户总结完成的工作）

如果选择TASK_COMPLETE，请先向用户说明已完成的操作和结果，然后再返回TASK_COMPLETE。

请只回复上述选项之一。
`;

        try {
            // 官方方式：使用完整的消息历史 + 决策提示
            const decisionModel = new ChatOpenAI({
                apiKey: OPENAI_API_KEY,
                modelName: DECISION_MODEL,
                configuration: {
                    baseURL: OPENAI_BASE_URL
                },
                temperature: 0
            });

            // 将决策提示添加到现有对话历史中
            const decisionMessages = [...state.messages, new HumanMessage(decisionPrompt)];
            const response = await decisionModel.invoke(decisionMessages);
            const decision = typeof response.content === 'string' ? response.content.trim() : '';

            // 根据AI决策返回路由
            switch (decision) {
                case 'CONTINUE_SEARCH':
                case 'EXECUTE_CODE':
                    return "action";
                case 'TASK_COMPLETE':
                    return "__end__";
                default:
                    // 安全默认：避免无限循环
                    return "__end__";
            }
        } catch (error) {
            console.error('[AI Decision] 决策失败:', error);
            return "__end__";
        }
    }

            // Action 节点: 执行AI的决策和工具调用
    private async actionNode(state: AgentState): Promise<{ messages: BaseMessage[] }> {
        const response = await this.model.invoke(state.messages);

        // 如果模型决定调用工具
        if (response.tool_calls && response.tool_calls.length > 0) {
            const toolMessages: ToolMessage[] = await Promise.all(
                response.tool_calls.map(async (toolCall) => {
                    const tool = this.tools.find((t) => t.name === toolCall.name);
                    if (!tool) {
                        throw new Error(`Tool ${toolCall.name} not found`);
                    }
                    const output = await tool.invoke(toolCall.args);
                    return new ToolMessage({
                        tool_call_id: toolCall.id!,
                        content: typeof output === 'string' ? output : JSON.stringify(output),
                        name: toolCall.name
                    });
                })
            );
            return { messages: [response, ...toolMessages] };
        }

        return { messages: [response] };
    }



    // Check 节点: AI自主决策
    private checkNode(state: AgentState): { messages: [] } {
        return { messages: [] };
    }

    public async run(query: string) {
        const compiledGraph = this.app.compile();

                // 添加system prompt到初始状态
        const systemPrompt = `你是一个专业的SpreadJS助手。简洁友好地与用户交流，避免冗长的步骤列表。`;

        const initialState: AgentState = {
            messages: [
                new SystemMessage(systemPrompt),
                new HumanMessage(query)
            ]
        };

        const finalState = await compiledGraph.invoke(initialState, { recursionLimit: 10 });
        return finalState;
    }

    public async* executeStream(query: string, sessionId: string) {
        const compiledGraph = this.app.compile();

        // 添加system prompt到初始状态
        const systemPrompt = `你是一个专业的SpreadJS助手。

**核心原则：**
- 简洁友好地与用户交流
- 理解用户需求后直接开始执行
- 避免冗长的步骤列表和重复说明
- 在调用工具前简要说明你要做什么

**交流风格：**
- 第一次：简要确认理解，然后开始行动
- 后续：基于已有信息继续执行，无需重新规划

示例：
用户："添加表格"
你："好的，我来为您添加表格。让我搜索相关API..."
（而不是列出1、2、3步骤）`;

        const initialState: AgentState = {
            messages: [
                new SystemMessage(systemPrompt),
                new HumanMessage(query)
            ]
        };

        let stepCounter = 0;

        const eventStream = compiledGraph.streamEvents(initialState, {
            version: "v2" as const,
            recursionLimit: 10
        });

        for await (const event of eventStream) {
            stepCounter++;
            const stepId = `step_${sessionId}_${stepCounter}`;
            const timestamp = new Date().toISOString();

            // 只记录重要的流事件
            if (event.event === 'on_chat_model_stream' && event.data.chunk?.content?.trim()) {
                console.log(`[AI Stream] "${event.data.chunk.content.substring(0, 50)}${event.data.chunk.content.length > 50 ? '...' : ''}"`);
            }

            switch (event.event) {
                                case 'on_chat_model_stream':
                    // AI思考的流式输出 - 过滤决策token
                    if (event.data.chunk?.content) {
                        const content = event.data.chunk.content;

                        // 精确匹配决策token和其组成部分
                        const decisionTokens = [
                            'CONTINUE_SEARCH', 'EXECUTE_CODE', 'TASK_COMPLETE',
                            'EXEC', 'UTE', '_CODE', 'TASK', '_COMPLETE',
                            'SEARCH', 'CODE', 'COMPLETE'
                        ];

                        const isDecisionToken = decisionTokens.some(token =>
                            content.trim() === token || content.includes(token)
                        );

                        if (!isDecisionToken && content.trim().length > 0) {
                            yield {
                                type: 'agent_message',
                                data: {
                                    content: content,
                                    timestamp,
                                    stepId: stepId + '_ai_stream',
                                    sessionId: sessionId
                                }
                            };
                        }
                    }
                    break;

                case 'on_tool_start':
                    // 工具调用开始 - 解析工具参数
                    let toolArgs = event.data.input;
                    if (typeof toolArgs === 'object' && toolArgs.input && typeof toolArgs.input === 'string') {
                        try {
                            toolArgs = JSON.parse(toolArgs.input);
                        } catch (e) {
                            // 参数解析失败时使用原始参数
                        }
                    }

                    yield {
                        type: 'agent_step',
                        data: {
                            step: 'tools',
                            output: {
                                lastToolCall: {
                                    name: event.name,
                                    args: toolArgs
                                }
                            },
                            timestamp,
                            stepId: stepId + '_tool_start',
                            sessionId: sessionId
                        }
                    };
                    break;

                case 'on_tool_end':
                    // 工具调用完成
                    let output: any = {
                        toolResult: event.data.output,
                        toolName: event.name
                    };

                    // execute-spreadjs特殊处理
                    if (event.name === 'execute-spreadjs' && typeof event.data.output === 'string') {
                        try {
                            const parsedContent = JSON.parse(event.data.output);
                            if (parsedContent.execute && parsedContent.query) {
                                output.lastToolCall = {
                                    name: 'execute-spreadjs',
                                    args: {
                                        executionPackage: {
                                            executeCode: parsedContent.execute,
                                            queryCode: parsedContent.query,
                                            id: `exec_${Date.now()}`
                                        }
                                    }
                                };
                            }
                        } catch (e) {
                            // JSON解析失败，使用原始数据
                        }
                    }

                    yield {
                        type: 'agent_step',
                        data: {
                            step: 'tools',
                            output: output,
                            timestamp,
                            stepId: stepId + '_tool_end',
                            sessionId: sessionId
                        }
                    };
                    break;

                default:
                    // 忽略其他事件
                    break;
            }
        }
    }
}
